#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/netfilter_ipv4.h>
#include <linux/skbuff.h>
#include <linux/udp.h>
#include <linux/tcp.h>
#include <linux/ip.h>
#include <net/tcp.h>
#include <linux/version.h>

#if LINUX_VERSION_CODE < KERNEL_VERSION(2, 6, 25)
#define NF_INET_PRE_ROUTING     NF_IP_PRE_ROUTING
#define NF_INET_LOCAL_IN        NF_IP_LOCAL_IN
#define NF_INET_FORWARD         NF_IP_FORWARD
#define NF_INET_LOCAL_OUT       NF_IP_LOCAL_OUT
#define NF_INET_POST_ROUTING    NF_IP_POST_ROUTING
#define NF_INET_NUMHOOKS        NF_IP_NUMHOOKS
#endif

#if LINUX_VERSION_CODE < KERNEL_VERSION(2,6,36)
#define tcp_flag_byte(th) (((u_int8_t *)th)[13])

#define TCPHDR_FIN 0x01
#define TCPHDR_SYN 0x02
#define TCPHDR_RST 0x04
#define TCPHDR_PSH 0x08
#define TCPHDR_ACK 0x10
#define TCPHDR_URG 0x20
#define TCPHDR_ECE 0x40
#define TCPHDR_CWR 0x80

#define TCPHDR_SYN_ECN	(TCPHDR_SYN | TCPHDR_ECE | TCPHDR_CWR)
#endif

static unsigned short g_mon_ports[] = { 7777, 7778 };

MODULE_LICENSE("GPL");
MODULE_AUTHOR("August");
MODULE_DESCRIPTION("A simple Linux port monitor");
MODULE_VERSION("0.1");

static bool check_monitor_port(unsigned short sport, unsigned short dport)
{
  int i = 0;

  for (; i < sizeof(g_mon_ports)/sizeof(unsigned short); i++)
  {
    if(sport == g_mon_ports[i] || dport == g_mon_ports[i])
    {
      return true;
    }
  }
  return false;
}

static void print_skb(bool in, struct sk_buff *skb, struct iphdr *iph, struct tcphdr *th)
{
  unsigned int sip = ntohl(iph->saddr);
  unsigned int dip = ntohl(iph->daddr);

  printk(KERN_INFO " %s source=%u.%u.%u.%u:%u dest=%u.%u.%u.%u:%u flag=0x%x seq=%u ack_seq=%u  datalen=%u\n", in ? ">>" : "<<",
      (sip >> 24) & 0xff, (sip >> 16) & 0xff, (sip >> 8) & 0xff, sip & 0xff, ntohs(th->source),
      (dip >> 24) & 0xff, (dip >> 16) & 0xff, (dip >> 8) & 0xff, dip & 0xff, ntohs(th->dest),
      tcp_flag_byte(th), ntohl(th->seq), ntohl(th->ack_seq), skb->len - (iph->ihl << 2) - (th->doff << 2));
}

#if LINUX_VERSION_CODE <= KERNEL_VERSION(2,6,20)
static unsigned int hook_func_in(unsigned int hooknum, struct sk_buff **pskb, const struct net_device *in, const struct net_device *out, int (*okfn)(struct sk_buff *))
#elif LINUX_VERSION_CODE < KERNEL_VERSION(3,13,0)
static unsigned int hook_func_in(unsigned int hooknum, struct sk_buff *skb, const struct net_device *in, const struct net_device *out, int (*okfn)(struct sk_buff *))
#elif LINUX_VERSION_CODE < KERNEL_VERSION(4,1,0)
static unsigned int hook_func_in(const struct nf_hook_ops *ops, struct sk_buff *skb, const struct net_device *in, const struct net_device *out, int (*okfn)(struct sk_buff *))
#elif LINUX_VERSION_CODE < KERNEL_VERSION(4,4,0)
static unsigned int hook_func_in(const struct nf_hook_ops *ops, struct sk_buff *skb, const struct nf_hook_state *state)
#else
static unsigned int hook_func_in(void *priv, struct sk_buff *skb, const struct nf_hook_state *state)
#endif
{
  struct iphdr *ip_header;
  struct udphdr *udp_header;
  struct tcphdr *tcp_header;

#if LINUX_VERSION_CODE <= KERNEL_VERSION(2,6,20)
    struct sk_buff *skb;
    if(NULL == pskb)
    {
        return NF_ACCEPT;
    }
    skb = *pskb;
    ip_header = skb->nh.iph;
#else
    if(NULL == skb)
    {
        return NF_ACCEPT;
    }
    ip_header = (struct iphdr *)skb_network_header(skb);
#endif
  switch (ip_header->protocol)
  {
  case IPPROTO_UDP:
    udp_header = (struct udphdr *)skb_transport_header(skb);
    break;
  case IPPROTO_TCP:
    //CentOS6.5  skb_transport_header has bug
    //tcp_header = (struct tcphdr *)skb_transport_header(skb);
    tcp_header = (struct tcphdr *)((unsigned char *)ip_header + ip_header->ihl * 4); 
    if (check_monitor_port(ntohs(tcp_header->source), ntohs(tcp_header->dest)))
    {
      print_skb(true, skb, ip_header, tcp_header);
    }
  }
  return NF_ACCEPT;
}

#if LINUX_VERSION_CODE <= KERNEL_VERSION(2,6,20)
static unsigned int hook_func_out(unsigned int hooknum, struct sk_buff **pskb, const struct net_device *in, const struct net_device *out, int (*okfn)(struct sk_buff *))
#elif LINUX_VERSION_CODE < KERNEL_VERSION(3,13,0)
static unsigned int hook_func_out(unsigned int hooknum, struct sk_buff *skb, const struct net_device *in, const struct net_device *out, int (*okfn)(struct sk_buff *))
#elif LINUX_VERSION_CODE < KERNEL_VERSION(4,1,0)
static unsigned int hook_func_out(const struct nf_hook_ops *ops, struct sk_buff *skb, const struct net_device *in, const struct net_device *out, int (*okfn)(struct sk_buff *))
#elif LINUX_VERSION_CODE < KERNEL_VERSION(4,4,0)
static unsigned int hook_func_out(const struct nf_hook_ops *ops, struct sk_buff *skb, const struct nf_hook_state *state)
#else
static unsigned int hook_func_out(void *priv, struct sk_buff *skb, const struct nf_hook_state *state)
#endif
{
  struct iphdr *ip_header;
  struct udphdr *udp_header;
  struct tcphdr *tcp_header;

#if LINUX_VERSION_CODE <= KERNEL_VERSION(2,6,20)
    struct sk_buff *skb;
    if(NULL == pskb)
    {
        return NF_ACCEPT;
    }
    skb = *pskb;
    ip_header = skb->nh.iph;
#else
    if(NULL == skb)
    {
        return NF_ACCEPT;
    }
    ip_header = (struct iphdr *)skb_network_header(skb);
#endif
  switch (ip_header->protocol)
  {
  case IPPROTO_UDP:
    udp_header = (struct udphdr *)skb_transport_header(skb);
    break;
  case IPPROTO_TCP:
    //tcp_header = (struct tcphdr *)skb_transport_header(skb);
    tcp_header = (struct tcphdr *)((unsigned char *)ip_header + ip_header->ihl * 4); 
    if (check_monitor_port(ntohs(tcp_header->source), ntohs(tcp_header->dest)))
    {
      print_skb(false, skb, ip_header, tcp_header);
    }
  }
  return NF_ACCEPT;
}

static struct nf_hook_ops nfho[] = {
  {
    .hook = hook_func_in,              /* hook function */
    .hooknum = NF_INET_LOCAL_IN, /* watch all packets */
    .pf = PF_INET,                  /* ip protocol family */
    .priority = NF_IP_PRI_FIRST,    /* high priority */
  }
  ,
  {
    .hook = hook_func_out,              /* hook function */
    .hooknum = NF_INET_LOCAL_OUT, /* watch all packets */
    .pf = PF_INET,                  /* ip protocol family */
    .priority = NF_IP_PRI_FIRST,    /* high priority */
  }
};

static int __init init_nf(void) {
  printk(KERN_INFO "netmon init.\n");
#if LINUX_VERSION_CODE < KERNEL_VERSION(4, 4, 0)
  nf_register_hooks(nfho, ARRAY_SIZE(nfho));
#else
  nf_register_net_hooks(&init_net, nfho, ARRAY_SIZE(nfho));
#endif
  return 0;
}

static void __exit exit_nf(void) {
  printk(KERN_INFO "netmon exit.\n");
#if LINUX_VERSION_CODE < KERNEL_VERSION(4, 4, 0)
  nf_unregister_hooks(nfho, ARRAY_SIZE(nfho));
#else
  nf_unregister_net_hooks(&init_net, nfho, ARRAY_SIZE(nfho));
#endif
}

module_init(init_nf);
module_exit(exit_nf);
